from functools import reduce
from json import load
from math import cos, floor, pi, sin
from os import listdir, mkdir, remove, rmdir
from os.path import dirname, isdir, join
from typing import Callable

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from docx.shared import Cm
from docxtpl import DocxTemplate, InlineImage
from scipy import integrate, interpolate, linalg, misc

deg = pi / 180
# массив значений угла фи для построения графиков функций
fi_arr = np.array([i * deg for i in range(362)])
# круг единичного радиуса для построения графиков
ring = np.array([[cos(fi), sin(fi)] for fi in fi_arr]).T
quad = integrate.quad
context = {}


def polar_to_xy(pt: [float, float], scale: float) -> (float, float):
    """Преобразует точку из полярных координат в ХУ с учетом масштаба

    Parameters
    ----------
    pt : [float,float]
        Исходная точка в полярных координатах
    float : float
        Масштаб графика

    Returns
    -------
    (float,float)
        Точка в координатах ХУ
    """
    return np.array(
        [(1 + scale * pt[1]) * cos(pt[0]), (1 + scale * pt[1]) * sin(pt[0])]
    )


def plot_function(
    func: Callable[[float], float], plotLabel: str, units: str, fileName: str
) -> (mpl.figure, mpl.axes):
    """Сохраняет форматированный график функции в файл.

    Parameters
    ----------
    func : {[Callable[[float], float]]}
        _description_
    mpl : _type_
        _description_
    """
    # массив значений функции
    values = np.array([func(fi) for fi in fi_arr])
    max_point = (0, values[0])  # координаты максимума (fi,r)
    min_point = (0, values[0])  # координаты минимума (fi,r)
    for i in range(len(values)):
        if values[i] > max_point[1]:
            max_point = (fi_arr[i], values[i])
        if values[i] < min_point[1]:
            min_point = (fi_arr[i], values[i])
    # масштаб построения
    scale = 0.5 / max(abs(max_point[1]), abs(min_point[1]), 0.0001)
    # values в XY  с учетом scale
    funcPoints = [
        (1 + values * scale) * np.cos(fi_arr),
        (1 + values * scale) * np.sin(fi_arr),
    ]
    canv, plot = plt.subplots(figsize=(7, 7))
    plot.set_xlim(xmin=-2, xmax=2)
    plot.set_ylim(ymin=-2, ymax=2)
    max_xy = polar_to_xy(max_point, scale)
    min_xy = polar_to_xy(min_point, scale)
    (line1,) = plot.plot(*ring, label="Ось", color="black")
    (line2,) = plot.plot(*funcPoints, label=plotLabel, color="blue")
    plot.plot(*max_xy, "o", color="red")
    plot.plot(*min_xy, "o", color="red")
    plot.legend(handles=[line1, line2], loc="upper left")
    plot.annotate(
        f"{max_point[1]:.2f} {units}",
        xy=max_xy,
        xytext=max_xy + np.array([0, 0.05]),
        ha="center",
        va="bottom",
        bbox={"boxstyle": "round", "fc": "w"},
    )
    plot.annotate(
        f"{min_point[1]:.2f} {units}",
        xy=min_xy,
        xytext=min_xy + np.array([0, 0.05]),
        ha="center",
        va="bottom",
        bbox={"boxstyle": "round", "fc": "w"},
    )
    canv.savefig(fileName)


def linterp(func: Callable[[float], float]) -> Callable[[float], float]:
    """Возвращает интерполяцию функции на интервале 0 2*pi"""
    f_pts = [func(fi) for fi in fi_arr]
    interp = interpolate.interp1d(fi_arr, f_pts)
    return lambda fi: interp(fi - 2 * pi * floor(fi / (2 * pi)))


# %%
# создание временной директории для хранения файлов
ABS_PATH = dirname(__file__)
TMP_PATH = join(ABS_PATH, "tmp")
TEMPLATE_PATH = join(ABS_PATH, "ring_template.docx")
OUT_PATH = join(ABS_PATH, "report.docx")
if not isdir(TMP_PATH):
    mkdir(TMP_PATH)


# %%
# загрузка исходных данных (файл src.json)
with open(join(ABS_PATH, "src.json"), "r") as f:
    src = load(f)


# %%
# вычисление строковых значений
for key, val in src.items():
    if key != "loads":
        src[key] = eval(str(val))

Q_MAX, P_MAX, M_MAX = [src["Q_MAX"], src["P_MAX"], src["M_MAX"]]
context["srcE"] = src["E"] / 1e6
context["srcR"] = src["R"]
context["qValue"] = Q_MAX
context["mValue"] = M_MAX
context["pValue"] = P_MAX


load_q = []  # Распределенные нагрузки
load_f_radial = []  # Сосредоточенные радиальные силы
load_m = []  # Сосредоточенные моменты
for item in src["loads"]:
    match item:
        case {"type": "Fr", "pos": pos, "value": value}:
            item["pos"] = eval(str(pos))
            item["value"] = eval(str(value))
            load_f_radial.append(item)

        case {"type": "q", "begin": begin, "end": end}:
            item["begin"] = eval(str(begin))
            item["end"] = eval(str(end))
            load_q.append(item)

        case {"type": "M", "pos": pos, "value": value}:
            item["pos"] = eval(str(pos))
            item["value"] = eval(str(value))
            load_m.append(item)


# %%
# построение функции распределенной нагрузки
q_list = list(
    map(
        lambda entry: (
            lambda fi: eval(entry["value"])
            if fi >= entry["begin"] and fi <= entry["end"]
            else 0
        ),
        load_q,
    )
)


# %%
# Задание суммарной распределенной нагрузки
def q_radial(fi: float) -> float:
    """Радиальная нагрузка в зависимости от аргумента"""
    return reduce(lambda acc, f: acc + f(fi), q_list, 0)


# %%
# Суммарная внешняя нагрузка

ext_fx = quad(lambda fi: q_radial(fi) * cos(fi), 0, 2 * pi)[0] + reduce(
    lambda acc, item: item["value"] * cos(item["pos"]) + acc,
    load_f_radial,
    0,
)
ext_fy = quad(lambda fi: q_radial(fi) * sin(fi), 0, 2 * pi)[0] + reduce(
    lambda acc, item: item["value"] * sin(item["pos"]) + acc,
    load_f_radial,
    0,
)
ext_m = reduce(lambda acc, item: acc + item["value"], load_m, 0)

# %%
# Вычисление уравновешивающей распределенной нагрузки qt = A + B*cos(fi) + C*sin(fi)

a = np.array(
    [
        [
            quad(lambda fi: cos(fi + pi / 2), 0, 2 * pi)[0],
            quad(lambda fi: cos(fi) * cos(fi + pi / 2), 0, 2 * pi)[0],
            quad(lambda fi: sin(fi) * cos(fi + pi / 2), 0, 2 * pi)[0],
        ],
        [
            quad(lambda fi: sin(fi + pi / 2), 0, 2 * pi)[0],
            quad(lambda fi: cos(fi) * sin(fi + pi / 2), 0, 2 * pi)[0],
            quad(lambda fi: sin(fi) * sin(fi + pi / 2), 0, 2 * pi)[0],
        ],
        [
            2 * pi * src["R"],
            0,
            0,
        ],
    ]
)
b = np.array([-ext_fx, -ext_fy, -ext_m])
coefs = linalg.solve(a, b)
context["coef_A"] = f"{coefs[0]:.2f}"
context["coef_B"] = f"{coefs[1]:.2f}"
context["coef_C"] = f"{coefs[2]:.2f}"


def q_t(fi: float) -> float:
    """Уравновешивающая распределенная нагрузка"""
    return coefs[0] + coefs[1] * cos(fi) + coefs[2] * sin(fi)


# %%
# Внутренние силовые факторы от внешней нагрузки


def h(x: float) -> float:
    """Функция Хевисайда"""
    return 1 if x > 0 else 0


r = src["R"]


@linterp
def n_ext(fi: float) -> float:
    """Нормальная сила в сечении от внешней нагрузки"""
    val = 0
    for force in load_f_radial:
        val += h(fi - force["pos"]) * force["value"] * cos(pi / 2 - fi + force["pos"])
    return (
        val
        + quad(
            lambda a: q_radial(fi - a) * r * sin(a) + q_t(fi - a) * r * cos(pi + a),
            0,
            fi,
            epsabs=0.0001,
            limit=200,
        )[0]
    )


@linterp
def q_ext(fi: float) -> float:
    """Перерезывающая сила в сечении от внешней нагрузки"""
    val = 0
    for force in load_f_radial:
        val += h(fi - force["pos"]) * force["value"] * cos(fi - pi - force["pos"])
    return (
        val
        + quad(
            lambda a: q_radial(fi - a) * r * cos(a + pi) - q_t(fi - a) * r * sin(a),
            0,
            fi,
            epsabs=0.0001,
            limit=200,
        )[0]
    )


@linterp
def m_ext(fi: float) -> float:
    """Изгибающий момент в сечении от внешней нагрузки"""
    val = 0
    for force in load_m:
        val += h(fi - force["pos"]) * force["value"]
    for force in load_f_radial:
        val += h(fi - force["pos"]) * force["value"] * r * sin(fi - force["pos"])
    return (
        val
        + quad(
            lambda a: q_radial(fi - a) * r * r * sin(a)
            + q_t(fi - a) * r * r * (1 - cos(a)),
            0,
            fi,
            epsabs=0.0001,
            limit=200,
        )[0]
    )


# %%
# Раскрытие статической неопределимости

EJ = src["E"] * (src["b"] ** 3) * src["h"] / 12

# моменты от единичных сил
M = [lambda fi: r * (1 - cos(fi)), lambda fi: r * sin(fi), lambda fi: 1]
# матрица коэффициентов
a = np.array(
    [
        [
            quad(lambda fi: M[i](fi) * M[j](fi), 0, 2 * pi, epsabs=0.0001)[0] / EJ
            for j in range(3)
        ]
        for i in range(3)
    ]
)
# Вектор правой части
b = np.array(
    [
        -quad(
            lambda fi: M[i](fi) * m_ext(fi), 0, 2 * pi, epsabs=0.0001, full_output=True
        )[0]
        / EJ
        for i in range(3)
    ]
)
x = linalg.solve(a, b)
context["x1"] = f"{x[0]:.2f}"
context["x2"] = f"{x[1]:.2f}"
context["x3"] = f"{x[2]:.2f}"


def m_sum(fi: float) -> float:
    """Результирующий момент в сечении"""
    return m_ext(fi) + x[0] * M[0](fi) + x[1] * M[1](fi) + x[2] * M[2](fi)


def n_sum(fi: float) -> float:
    """Результирующая осевая сила в сечении"""
    return n_ext(fi) - x[0] * cos(fi) + x[1] * sin(fi)


def q_sum(fi: float) -> float:
    """Результирующая перерезывающая сила в сечении"""
    return q_ext(fi) + x[0] * sin(fi) - x[1] * cos(fi)


# %%
# Определение перемещений
def dw(fi: float, y: [float]) -> [float]:
    """
    Функция, определяющая правую часть ДУ: w'' = -w + M * R^2 / (EJ).

    * fi - текущий угол
    * y = [w',w] - угол поворота сечения кольца
    """
    return [-y[1] - m_sum(fi) * r * r / EJ, y[0]]


sol = integrate.solve_ivp(
    dw, (0, 365 * deg), [0, 0], dense_output=True, t_eval=fi_arr, max_step=1e-3
)
w0 = interpolate.interp1d(sol.t, sol.y[1])  # частное решение ДУ


@linterp
def v0(fi: float) -> float:
    """Частное решение уравнения W = - V'"""
    return quad(w0, 0, fi, full_output=True)[0]


# определение констант в общем решении ДУ
C0 = (-0.5 / pi) * quad(v0, 0, 2 * pi, full_output=True)[0]
C1 = (-1 / pi) * quad(lambda fi: v0(fi) * cos(fi), 0, 2 * pi, full_output=True)[0]
C2 = (-1 / pi) * quad(lambda fi: v0(fi) * sin(fi), 0, 2 * pi, full_output=True)[0]

context["c1"] = f"{C0:.5f}"
context["c2"] = f"{C1:.5f}"
context["c3"] = f"{C2:.5f}"


def v(fi: float) -> float:
    return v0(fi) + C0 + C1 * cos(fi) + C2 * sin(fi)


def w(fi: float) -> float:
    return -misc.derivative(v, fi, 1e-2, n=1)


def q(fi: float) -> float:
    return (-1 / r) * (misc.derivative(w, fi, 1e-2) + v(fi))


# %%
# Построение графиков функций

plot_function(q_t, "qt(φ)", "Н/м", join(TMP_PATH, "qT.png"))
plot_function(q_radial, "qr(φ)", "Н/м", join(TMP_PATH, "qR.png"))
plot_function(n_ext, "N(φ)", "Н", join(TMP_PATH, "Next.png"))
plot_function(q_ext, "Q(φ)", "Н", join(TMP_PATH, "Qext.png"))
plot_function(m_ext, "М(φ)", "Н*м", join(TMP_PATH, "Mext.png"))
plot_function(m_sum, "M(φ)", "Н*м", join(TMP_PATH, "Msum.png"))
plot_function(n_sum, "N(φ)", "Н", join(TMP_PATH, "Nsum.png"))
plot_function(q_sum, "Q(φ)", "Н", join(TMP_PATH, "Qsum.png"))
plot_function(lambda fi: v(fi) * 1000, "v(φ)", "мм", join(TMP_PATH, "v.png"))
plot_function(lambda fi: w(fi) * 1000, "w(φ)", "мм", join(TMP_PATH, "w.png"))
plot_function(lambda fi: q(fi) / deg, "q(φ)", "°", join(TMP_PATH, "q.png"))

doc = DocxTemplate(TEMPLATE_PATH)
context["img_qr"] = InlineImage(doc, join(TMP_PATH, "qR.png"), Cm(8), Cm(8))
context["img_qt"] = InlineImage(doc, join(TMP_PATH, "qT.png"), Cm(8), Cm(8))
context["img_Next"] = InlineImage(doc, join(TMP_PATH, "Next.png"), Cm(8), Cm(8))
context["img_Qext"] = InlineImage(doc, join(TMP_PATH, "Qext.png"), Cm(8), Cm(8))
context["img_Mext"] = InlineImage(doc, join(TMP_PATH, "Mext.png"), Cm(8), Cm(8))
context["img_Nsum"] = InlineImage(doc, join(TMP_PATH, "Nsum.png"), Cm(8), Cm(8))
context["img_Qsum"] = InlineImage(doc, join(TMP_PATH, "Qsum.png"), Cm(8), Cm(8))
context["img_Msum"] = InlineImage(doc, join(TMP_PATH, "Msum.png"), Cm(8), Cm(8))
context["img_v"] = InlineImage(doc, join(TMP_PATH, "v.png"), Cm(8), Cm(8))
context["img_w"] = InlineImage(doc, join(TMP_PATH, "w.png"), Cm(8), Cm(8))
context["img_q"] = InlineImage(doc, join(TMP_PATH, "q.png"), Cm(8), Cm(8))
doc.render(context)
doc.save(OUT_PATH)


# %%
# удаление временной директории
for fn in listdir(TMP_PATH):
    remove(join(TMP_PATH, fn))
rmdir(TMP_PATH)
